#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
用户标签生成模块
用于为广电用户数据生成标签并转换为宽表格式
"""
import pandas as pd
import numpy as np
import random
import os

def generate_user_labels(df, rule_based_ratio=0.3):
    """
    为用户生成标签
    
    参数:
    df: 处理后的用户数据框
    rule_based_ratio: 基于规则打标的比例，默认为30%
    
    返回:
    带标签的用户数据框
    """
    labeled_df = df.copy()
    
    # 初始化标签列
    labeled_df['user_label'] = None
    labeled_df['label_type'] = None  # 'rule'表示规则标签，'random'表示随机标签
    
    # 标签类别列表
    label_categories = ['高价值用户', '忠诚用户', '新用户', '潜在流失用户', '一般用户']
    
    # 确定基于规则和随机打标的用户数量
    total_users = len(labeled_df)
    rule_based_count = int(total_users * rule_based_ratio)
    
    # 获取所有用户索引
    all_indices = list(range(total_users))
    
    # 随机选择一部分用户用于基于规则打标
    rule_based_indices = random.sample(all_indices, rule_based_count)
    random_indices = [i for i in all_indices if i not in rule_based_indices]
    
    # 1. 基于规则打标
    for idx in rule_based_indices:
        # 这里可以根据实际数据特征制定规则
        # 示例规则：
        # - 有多个产品的用户标记为高价值用户
        # - 最近活跃的用户标记为忠诚用户
        # - 新注册的用户标记为新用户
        # - 长期不活跃的用户标记为潜在流失用户
        
        # 模拟规则逻辑（实际应用中需要根据具体字段调整）
        user_row = labeled_df.iloc[idx]
        
        # 假设如果有'sm_name'列，可以基于产品类型判断
        if 'sm_name' in user_row.index and isinstance(user_row['sm_name'], str):
            if '数字电视' in user_row['sm_name'] and '宽带' in user_row['sm_name']:
                labeled_df.at[idx, 'user_label'] = '高价值用户'
                labeled_df.at[idx, 'label_type'] = 'rule'
                continue
        
        # 假设如果有'run_time'列，可以基于活跃时间判断
        if 'run_time' in user_row.index:
            try:
                # 转换为datetime
                run_time = pd.to_datetime(user_row['run_time'])
                # 如果是最近一个月的记录
                if (pd.Timestamp.now() - run_time).days < 30:
                    labeled_df.at[idx, 'user_label'] = '忠诚用户'
                    labeled_df.at[idx, 'label_type'] = 'rule'
                    continue
                # 如果是超过6个月的记录
                elif (pd.Timestamp.now() - run_time).days > 180:
                    labeled_df.at[idx, 'user_label'] = '潜在流失用户'
                    labeled_df.at[idx, 'label_type'] = 'rule'
                    continue
            except:
                pass
        
        # 如果没有适用的规则，则分配一个随机标签
        labeled_df.at[idx, 'user_label'] = random.choice(label_categories)
        labeled_df.at[idx, 'label_type'] = 'random_rule_fallback'
    
    # 2. 随机打标剩余用户
    for idx in random_indices:
        labeled_df.at[idx, 'user_label'] = random.choice(label_categories)
        labeled_df.at[idx, 'label_type'] = 'random'
    
    print(f"生成标签完成：")
    print(f"- 基于规则打标: {rule_based_count} 条")
    print(f"- 随机打标: {len(random_indices)} 条")
    print(f"标签分布：")
    print(labeled_df['user_label'].value_counts())
    
    return labeled_df

def convert_to_wide_format(df, id_column='phone_no', label_column='user_label'):
    """
    将数据转换为宽表格式
    
    参数:
    df: 带标签的数据框
    id_column: 唯一标识列名
    label_column: 标签列名
    
    返回:
    宽表格式的数据框
    """
    # 创建标签的独热编码
    if label_column in df.columns:
        # 获取所有唯一标签
        unique_labels = df[label_column].unique()
        
        # 创建宽表
        wide_df = df.copy()
        
        # 为每个标签创建二值列
        for label in unique_labels:
            wide_df[f'label_{label}'] = (wide_df[label_column] == label).astype(int)
        
        print(f"转换为宽表格式完成")
        print(f"新增 {len(unique_labels)} 个标签列")
        
        return wide_df
    else:
        print(f"警告: 标签列 '{label_column}' 不存在")
        return df

def process_and_label_data(input_file, output_file=None, rule_based_ratio=0.3):
    """
    处理数据并添加标签，转换为宽表格式
    
    参数:
    input_file: 输入文件路径
    output_file: 输出文件路径
    rule_based_ratio: 基于规则打标的比例
    
    返回:
    处理后的宽表数据框
    """
    try:
        # 读取数据
        print(f"读取文件: {input_file}")
        df = pd.read_csv(input_file, sep=';', encoding='utf-8')
        
        # 生成标签
        labeled_df = generate_user_labels(df, rule_based_ratio)
        
        # 转换为宽表
        wide_df = convert_to_wide_format(labeled_df)
        
        # 保存结果
        if output_file:
            wide_df.to_csv(output_file, sep=';', index=False, encoding='utf-8')
            print(f"结果已保存到: {output_file}")
        
        return wide_df
        
    except Exception as e:
        print(f"处理过程中出错: {str(e)}")
        import traceback
        traceback.print_exc()
        raise

def batch_process_and_label(input_dir, output_dir, rule_based_ratio=0.3):
    """
    批量处理目录中的所有文件并添加标签
    
    参数:
    input_dir: 输入目录
    output_dir: 输出目录
    rule_based_ratio: 基于规则打标的比例
    """
    import glob
    
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)
    
    # 查找所有CSV文件
    csv_files = glob.glob(os.path.join(input_dir, "*.csv"))
    
    if not csv_files:
        print(f"在目录 {input_dir} 中未找到CSV文件")
        return
    
    print(f"找到 {len(csv_files)} 个CSV文件")
    
    for csv_file in csv_files:
        print(f"\n处理文件: {os.path.basename(csv_file)}")
        
        # 生成输出文件名
        base_name = os.path.basename(csv_file)
        output_file = os.path.join(output_dir, f"labeled_{base_name}")
        
        # 处理并添加标签
        process_and_label_data(csv_file, output_file, rule_based_ratio)
    
    print(f"\n批量处理完成")

# 主函数演示使用方法
def main():
    """
    主函数，演示使用方法
    """
    print("用户标签生成模块使用示例")
    print("="*50)
    
    # 这里可以添加示例代码，展示如何使用这些函数
    print("\n使用方法示例:")
    print("1. 生成用户标签:")
    print("   df = pd.read_csv('data.csv', sep=';')")
    print("   labeled_df = generate_user_labels(df)")
    
    print("\n2. 转换为宽表:")
    print("   wide_df = convert_to_wide_format(labeled_df)")
    
    print("\n3. 一站式处理:")
    print("   process_and_label_data('input.csv', 'output.csv')")
    
    print("\n4. 批量处理:")
    print("   batch_process_and_label('./input_dir', './output_dir')")

if __name__ == "__main__":
    main()